from typing import List

import pandas as pd
from more_itertools import chunked
from pydantic import BaseModel, field_validator
from tqdm import tqdm
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams


class LLMGenerationOutput(BaseModel):
    text: str
    label: List[str] = []
    explanation: str = ""

    @field_validator("label")
    @classmethod
    def validate_label(cls, label):
        return [
            (
                item
                if item
                in ["家人", "朋友", "同事", "上级", "电商客服", "老师", "伴侣", "其他"]
                else "未知"
            )
            for item in label
        ]

    @classmethod
    def from_str(cls, s: str):
        try:
            # find ```json``` block
            start = s.find("{")
            end = s.rfind("}") + 1
            obj = cls.model_validate_json(s[start:end])

            explanation = s[:start].strip().replace("分析：", "")
            obj.explanation = explanation
        except:
            obj = cls(text=s, label=["未知"], explanation="解析失败")
        return obj


# Common prefix.
prefix = """
你是一名数据标注员，给定一段用户在输入框中输入的内容，请判断一下你当前聊天对象的身份或者你们之间的关系。例如，你可能正在与以下类型的对象聊天：

家人：比如父母、兄弟姐妹或其他亲属。
朋友：私人生活中的好友或熟人。
同事：一起工作的团队成员或职场伙伴。
上级：你的直接领导或上级管理者。
电商客服：在线购物平台的客服人员。
老师：教育机构中的教员或导师。
伴侣：你的配偶或情侣，男女朋友。
其他：以上未列出的其他关系。

请根据你当前聊天的对象，选择最符合的选项。你的回答将有助于我们更好地理解你的聊天背景，并为你提供更精准的服务。
如果你不确定，可以选择多个候选选项，最有可能的选项请排在前面。
如果上述类别都不满足，可以选择其他类型。

请先对输入内容进行分析，再选择最符合的选项。选项使用以下json格式输出:

分析：xxx
{
    "label": ["xxx"],
    "text": "xxx"
}
"""

dataframe = pd.read_excel("./data/role_cls/predict.xlsx", dtype=str)
if "角色标签" not in dataframe.columns:
    dataframe["角色标签"] = None
# select 输入 is not null and 角色标签 is null
dataframe = dataframe[dataframe["输入"].notnull() & dataframe["角色标签"].isnull()]
# dataframe = dataframe.iloc[:100]


# Sample prompts.
prompts = dataframe["输入"].tolist()
# tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-57B-A14B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B-Instruct")

generating_prompts = []

for prompt in prompts:
    messages = [
        {"role": "system", "content": prefix},
        {"role": "user", "content": prompt},
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    generating_prompts.append(text)

# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.0, max_tokens=128)

# Increase gpu_memory_utilization to 0.9
prefix_cached_llm = LLM(
    model="Qwen/Qwen2-7B-Instruct",
    enable_prefix_caching=True,
    gpu_memory_utilization=0.95,
    max_model_len=2048,
)

# Warmup so that the shared prompt's KV cache is computed.
prefix_cached_llm.generate(generating_prompts[0], sampling_params, use_tqdm=False)

# Generate with prefix caching.
outputs = []
for batch in tqdm(chunked(generating_prompts, 8), total=len(generating_prompts) // 8):
    outputs += prefix_cached_llm.generate(batch, sampling_params, use_tqdm=False)

labels = []
# Print the outputs. You should see the same outputs as before.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    item = LLMGenerationOutput.from_str(generated_text)
    labels.append(item.label)


dataframe["角色标签"] = [",".join(label) for label in labels]
dataframe.to_excel("./data/role_cls/predict_with_label.xlsx", index=False)
